package com.free.mini.gateway.filter;

import com.free.mini.gateway.constants.GatewayConstant;
import org.reactivestreams.Publisher;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.cloud.gateway.filter.GatewayFilterChain;
import org.springframework.cloud.gateway.filter.GlobalFilter;
import org.springframework.cloud.gateway.filter.factory.rewrite.CachedBodyOutputMessage;
import org.springframework.cloud.gateway.support.BodyInserterContext;
import org.springframework.core.annotation.Order;
import org.springframework.core.io.buffer.DataBuffer;
import org.springframework.core.io.buffer.DataBufferUtils;
import org.springframework.http.HttpHeaders;
import org.springframework.http.MediaType;
import org.springframework.http.server.reactive.ServerHttpRequest;
import org.springframework.http.server.reactive.ServerHttpRequestDecorator;
import org.springframework.stereotype.Component;
import org.springframework.web.reactive.function.BodyInserter;
import org.springframework.web.reactive.function.BodyInserters;
import org.springframework.web.reactive.function.server.HandlerStrategies;
import org.springframework.web.reactive.function.server.ServerRequest;
import org.springframework.web.server.ServerWebExchange;

import com.free.mini.gateway.function.DecryptFunction;

import lombok.extern.slf4j.Slf4j;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;

/**
 * @author dongsufeng
 * @version 4.0
 * @date 2023/8/9 12:16
 */
@Component
@Slf4j
@Order(1)
public class DecryptRequestFilter implements GlobalFilter {
    @Autowired
    private DecryptFunction decryptFunction;

    @Override
    public Mono<Void> filter(ServerWebExchange exchange, GatewayFilterChain chain) {
        ServerRequest serverRequest = ServerRequest.create(exchange, HandlerStrategies.withDefaults().messageReaders());
        GatewayContext gatewayContext = (GatewayContext) exchange.getAttributes().get(GatewayConstant.CACHE_GATEWAY_CONTEXT);
        if (!gatewayContext.getMethod().name().equalsIgnoreCase("post")){
            return chain.filter(exchange);
        }
        if (!gatewayContext.getContentType().getType().equals(MediaType.APPLICATION_JSON.getType())
            || !MediaType.APPLICATION_JSON.getSubtype().equals(gatewayContext.getContentType().getSubtype())){
            return chain.filter(exchange);
        }
        Publisher<String> modifiedBody = serverRequest.bodyToMono(String.class).flatMap((originalBody) -> {
            return (Mono)decryptFunction.apply(exchange, originalBody);
        }).switchIfEmpty(Mono.defer(() -> {
            return (Mono)decryptFunction.apply(exchange, null);
        }));
        BodyInserter bodyInserter = BodyInserters.fromPublisher(modifiedBody, String.class);
        HttpHeaders headers = new HttpHeaders();
        headers.putAll(exchange.getRequest().getHeaders());
        headers.remove("Content-Length");
        headers.set(HttpHeaders.CONTENT_TYPE,"application/json");
        CachedBodyOutputMessage outputMessage = new CachedBodyOutputMessage(exchange, headers);
        return bodyInserter.insert(outputMessage, new BodyInserterContext()).then(Mono.defer(() -> {
            ServerHttpRequest decorator = this.decorate(exchange, headers, outputMessage);
            return chain.filter(exchange.mutate().request(decorator).build());
        })).onErrorResume((throwable) -> {
            return this.release(exchange, outputMessage, (Throwable) throwable);
        });
    }
    protected Mono<Void> release(ServerWebExchange exchange, CachedBodyOutputMessage outputMessage, Throwable throwable) {
        return outputMessage.getBody().map(DataBufferUtils::release).then(Mono.error(throwable));
    }
    ServerHttpRequestDecorator decorate(ServerWebExchange exchange, HttpHeaders headers, CachedBodyOutputMessage outputMessage) {
        return new ServerHttpRequestDecorator(exchange.getRequest()) {
            @Override
            public HttpHeaders getHeaders() {
                long contentLength = headers.getContentLength();
                HttpHeaders httpHeaders = new HttpHeaders();
                httpHeaders.putAll(headers);
                if (contentLength > 0L) {
                    httpHeaders.setContentLength(contentLength);
                } else {
                    httpHeaders.set("Transfer-Encoding", "chunked");
                }

                return httpHeaders;
            }
            @Override
            public Flux<DataBuffer> getBody() {
                return outputMessage.getBody();
            }
        };
    }
}
